""" This runs the model once at the estimated parameters to save the results that
are later used to construct acceptance sets, fitted moment table, etc

"""

import pandas as pd
import sys
sys.path.append('./')

from src.models_new import simulation
from src.run_scripts import utils

# %% READ IN THE INPUTS -----------------------------------------------------------------
t = 'month'

data, delta, rho, c, weights = utils.read_in_data(
    path_moments_data='./models/smm_input/moments_empirical.csv',
    path_n_rigs=f"./models/first_stage/n_rigs",
    path_surplus_components="./models/surplus/surplus_components",
    path_surplus_grid='./models/surplus/surplus_grid_2_low_month.npy',
    path_df_state='./data_py/processed/states',
    path_delta='./models/smm_input/delta.csv',
    path_entry_cost='./models/smm_input/entry_cost.csv',
    path_rho='./models/smm_input/rho.csv',
    path_df_contracts='./data_py/processed/contracts_final.csv',
    path_price_match_values='./models/price_match/price_match_values',
    path_coefs_data='./models/smm_input/coefs_data',
    path_prob_match_predict_contracts='./models/robustness/prob_match_predict_contracts',
    path_prob_match_predict='./models/robustness/prob_match_predict',
    time_period=t,
    p_exit=1.0,
    use_myopic=True
)

#%% READ IN PARAMS ----------------------------------------------------------------------
params = pd.read_csv(
    './models/smm/params_smm_with_diff_new.csv', squeeze=True, index_col=[0]).to_dict()

# %% DO THE SIMULATION ------------------------------------------------------------------
args_config = {
    'params_fixed': {},
    'weights': weights,
    'mri_max': 2.15,
    'verbose': False,
    'verbose_output': True,
    'value_zero': False,
    'entry_prob_by_tau_by_ym': None,
    'tau_multiplier': 1,
    'x_names': list(params.keys()),
    'well_target': True
}
args_by_names = {**args_config, **data}

#%% DO SIMULATION -----------------------------------------------------------------------
results_benchmark = simulation.run_steps(
    x=list(params.values()),
    args_by_names=args_by_names
)

#%% GET MOMENTS INTO ONE OBJECT ---------------------------------------------------------
moments_all = pd.concat([
    pd.Series(results_benchmark['moments_sim']),
    results_benchmark['coefs_predicted']
])
moments_all['reneg'] = results_benchmark['mean_reneg']

#%% SAVE THE RESULTS --------------------------------------------------------------------
results_benchmark['cutoffs_min'].to_csv('./models/benchmark/cutoffs_min.csv')
results_benchmark['cutoffs_max'].to_csv('./models/benchmark/cutoffs_max.csv')
moments_all.to_csv('./models/benchmark/moments_simulated.csv')

#%% SAVE THE WELL PROB OF MATCHING ------------------------------------------------------
pd.DataFrame(
    results_benchmark['shares_by_state'],
    index=data['state_data']['date'],
    columns=['low', 'mid', 'high']
).to_csv('./models/benchmark/shares_by_state.csv')

